from PyQt6.QtWidgets import (QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QPushButton, QTextEdit, QComboBox, QLabel, QLineEdit,
                            QMessageBox, QInputDialog, QFileDialog)
from PyQt6.QtCore import Qt, pyqtSignal
from chat_display import ChatDisplay
from ai_models import OpenAIChat
from response_thread import AIResponseThread
from conversation_manager import ConversationManager
from conversation_manager_dialog import ConversationManagerDialog
from styles import COMBOBOX_STYLE, SETTINGS_BUTTON_STYLE, LABEL_STYLE, INPUT_STYLE, INPUT_TEXTEDIT_STYLE, SEND_BUTTON_STYLE
from model_manager import ModelManagerDialog
from toast import Toast

class InputManager(QTextEdit):
    message_sent = pyqtSignal()
    
    def __init__(self):
        super().__init__()
        self.init_ui()
        
    def init_ui(self):
        self.setFixedHeight(50)
        self.setPlaceholderText("输入消息...")
        self.setStyleSheet(INPUT_TEXTEDIT_STYLE)
        
    def keyPressEvent(self, event):
        # Alt+Enter 换行，Enter 发送
        if event.key() in (Qt.Key.Key_Return, Qt.Key.Key_Enter):
            if event.modifiers() == Qt.KeyboardModifier.AltModifier:
                # Alt+Enter 换行
                cursor = self.textCursor()
                cursor.insertText("\n")
            else:
                # Enter 发送
                self.send_message()
        else:
            super().keyPressEvent(event)
        
    def send_message(self):
        if self.toPlainText().strip():
            self.message_sent.emit()
        
    def get_input(self):
        """获取输入内容并清空输入框"""
        text = self.toPlainText().strip()
        self.clear()
        return text

class AIMainWindow(QMainWindow):
    def __init__(self, assistant):
        super().__init__()
        self.assistant = assistant
        self.settings_window = None
        self.openai_chat = OpenAIChat()
        self.openai_chat.model_changed.connect(self.on_model_changed)
        self.response_thread = None
        self.current_response = ""
        self.version = "1.3.0"
        
        self.initUI()
        
        # 创建对话管理器并连接信号
        self.conversation_manager = ConversationManager(self)
        self.conversation_manager.conversation_updated.connect(self.update_conversation_selector)
        self.conversation_manager.error_occurred.connect(self.chat_display.display_error)
        
        self.load_models()
        self.toast = Toast(self)

    def initUI(self):
        self.setWindowTitle(f"AI 界面 - 版本 {self.version}")
        self.setGeometry(100, 100, 330, 500)
        self.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
        
        # 修改窗口标志
        self.setWindowFlags(
            Qt.WindowType.Window |
            Qt.WindowType.WindowMinMaxButtonsHint |
            Qt.WindowType.WindowCloseButtonHint |
            Qt.WindowType.WindowStaysOnTopHint
        )
        
        # 创建中央窗口部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # 模型管理区域
        model_layout = QHBoxLayout()
        model_layout.setSpacing(10)  # 设置组件间距
        
        # 模型管理按钮
        self.manage_model_button = QPushButton("模型管理")
        self.manage_model_button.setFixedSize(100, 30)  # 固定按钮大小
        self.manage_model_button.setStyleSheet(SETTINGS_BUTTON_STYLE)
        self.manage_model_button.clicked.connect(self.open_model_manager)
        model_layout.addWidget(self.manage_model_button)
        
        # 模型选择下拉框
        self.model_selector = QComboBox()
        self.model_selector.setFixedHeight(30)  # 固定高度
        self.model_selector.setMinimumWidth(200)
        self.model_selector.setStyleSheet(COMBOBOX_STYLE)
        self.model_selector.currentTextChanged.connect(self.model_selector_changed)
        model_layout.addWidget(self.model_selector)
        
        # 添加关于按钮
        self.about_button = QPushButton("?")
        self.about_button.setFixedSize(30, 30)  # 与添加对话按钮大小一致
        self.about_button.setStyleSheet("""
            QPushButton {
                background-color: #007bff;
                border: none;
                border-radius: 15px;
                color: white;
                font-size: 16px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #0056b3;
            }
            QPushButton:pressed {
                background-color: #004085;
            }
        """)
        self.about_button.clicked.connect(self.show_about)
        model_layout.addWidget(self.about_button)
        
        model_layout.addStretch()
        layout.addLayout(model_layout)
        
        # 对话管理区域
        conversation_layout = QHBoxLayout()
        conversation_layout.setSpacing(10)  # 设置组件间距
        
        # 对话管理按钮
        self.manage_conversation_button = QPushButton("对话管理")
        self.manage_conversation_button.setFixedSize(100, 30)  # 固定按钮大小
        self.manage_conversation_button.setStyleSheet(SETTINGS_BUTTON_STYLE)
        self.manage_conversation_button.clicked.connect(self.open_conversation_manager)
        conversation_layout.addWidget(self.manage_conversation_button)
        
        # 对话选择下拉框
        self.conversation_selector = QComboBox()
        self.conversation_selector.setFixedHeight(30)  # 固定高度
        self.conversation_selector.setMinimumWidth(200)
        self.conversation_selector.setStyleSheet(COMBOBOX_STYLE)
        self.is_programmatic_change = False
        self.conversation_selector.currentIndexChanged.connect(
            lambda idx: self.conversation_changed(self.conversation_selector.itemText(idx))
        )
        conversation_layout.addWidget(self.conversation_selector)
        
        # 添加新对话按钮
        self.add_conversation_button = QPushButton("+")
        self.add_conversation_button.setFixedSize(30, 30)  # 正方形按钮
        self.add_conversation_button.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                border: none;
                border-radius: 15px;
                color: white;
                font-size: 20px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
            QPushButton:pressed {
                background-color: #3d8b40;
            }
        """)
        self.add_conversation_button.clicked.connect(self.create_new_conversation)
        conversation_layout.addWidget(self.add_conversation_button)
        
        conversation_layout.addStretch()
        layout.addLayout(conversation_layout)
        
        # 添加一些间距
        layout.addSpacing(10)
        
        # 添加聊天显示区域
        self.chat_display = ChatDisplay()
        self.chat_display.edit_message.connect(self.handle_edit_message)
        layout.addWidget(self.chat_display)
        
        # 创建底部布局
        bottom_layout = QHBoxLayout()
        
        # 添加输入区域
        self.input_manager = InputManager()
        self.input_manager.message_sent.connect(self.send_message)
        bottom_layout.addWidget(self.input_manager)
        
        # 添加发送按钮
        self.send_button = QPushButton("发送")
        self.send_button.setFixedSize(100, 40)
        self.send_button.setStyleSheet(SEND_BUTTON_STYLE)
        self.send_button.clicked.connect(self.input_manager.send_message)
        bottom_layout.addWidget(self.send_button)
        
        layout.addLayout(bottom_layout)

    def create_new_conversation(self):
        """创建新对话"""
        current_model = self.model_selector.currentText()
        conversation_id = self.conversation_manager.create_new_conversation(current_model)
        if conversation_id is not None:
            self.update_conversation_selector(current_model, conversation_id)
            self.chat_display.clear_messages()
            # 显示Toast提示
            self.toast.show_message("已创建新对话")

    def update_conversation_selector(self, model_name=None, select_conversation_id=None):
        """更新对话选择器"""
        if model_name is None:
            model_name = self.model_selector.currentText()
        
        self.is_programmatic_change = True  # 设置标志位
        self.conversation_selector.clear()
        
        conversations = self.conversation_manager.get_conversation_list(model_name)
        for conv_id, conv_name in conversations:
            self.conversation_selector.addItem(conv_name, conv_id)
            
        # 如果没有指定要选中的对话ID，则使用当前对话ID
        if select_conversation_id is None:
            select_conversation_id = self.conversation_manager.current_conversation.get(model_name)
            
        if select_conversation_id is not None:
            # 查找并选中指定的对话
            for i in range(self.conversation_selector.count()):
                if self.conversation_selector.itemData(i) == select_conversation_id:
                    self.conversation_selector.setCurrentIndex(i)
                    break
            
            # 清空显示区域
            self.chat_display.clear_messages()
            
            # 加载对话历史
            messages = self.conversation_manager.get_current_conversation_messages(model_name)
            for msg in messages:
                self.chat_display.add_message(msg["content"], role=msg["role"])
        
        self.is_programmatic_change = False  # 重置标志位

    def conversation_changed(self, conversation_name):
        """处理对话切换事件"""
        if self.is_programmatic_change:  # 如果是程序触发的改变，直接返回
            return
            
        current_model = self.model_selector.currentText()
        if not current_model or not conversation_name:
            return
            
        # 获取对话ID
        conversation_id = self.conversation_selector.currentData()
        if conversation_id is None:
            return
            
        # 设置当前对话
        if not self.conversation_manager.set_current_conversation(current_model, conversation_id):
            return
            
        # 清空显示区域
        self.chat_display.clear_messages()
        
        # 加载对话历史
        messages = self.conversation_manager.get_current_conversation_messages(current_model)
        for msg in messages:
            self.chat_display.add_message(msg["content"], role=msg["role"])

    def model_selector_changed(self, model_name):
        """处理模型切换事件"""
        if not model_name:
            return
            
        # 清空显示区域
        self.chat_display.clear_messages()
        
        # 更新对话选择器
        self.update_conversation_selector(model_name)
        
        # 如果模型没有对话，创建一个新对话
        conversations = self.conversation_manager.get_conversation_list(model_name)
        if not conversations:
            self.create_new_conversation()

    def send_message(self):
        message = self.input_manager.get_input()
        if message.strip():
            current_model = self.model_selector.currentText()
            if not current_model:
                self.chat_display.display_error("请先选择一个模型！")
                return

            # 确保使用当前选中的对话ID
            conversation_id = self.conversation_selector.currentData()
            if conversation_id is not None:
                self.conversation_manager.set_current_conversation(current_model, conversation_id)

            # 添加用户消息到显示和历史
            self.chat_display.add_message(message, role="user")
            self.conversation_manager.add_message(current_model, message, role="user")
            
            try:
                # 创建并启动响应线程
                self.response_thread = AIResponseThread(
                    self.openai_chat, 
                    message,
                    current_model,
                    self.conversation_manager.get_current_conversation_messages(current_model)
                )
                self.response_thread.response_ready.connect(self.handle_response)
                self.response_thread.start()
                
                self.send_button.setEnabled(False)
                self.send_button.setText("等待响应...")
                
            except Exception as e:
                self.chat_display.display_error(str(e))

    def handle_response(self, response):
        """处理AI响应"""
        try:
            if response.has_error():
                self.chat_display.display_error(response.error)
                self.send_button.setEnabled(True)
                self.send_button.setText("发送")
                return
            
            current_model = self.model_selector.currentText()
            
            if not response.is_final:
                self.chat_display.display_assistant_message(response.text, is_chunk=True)
            else:
                self.chat_display.display_assistant_message(response.text, is_chunk=False)
                # 将完整的AI响应添加到历史记录
                self.conversation_manager.add_message(current_model, response.text, role="assistant")
                self.send_button.setEnabled(True)
                self.send_button.setText("发送")
            
        except Exception as e:
            self.chat_display.display_error(f"显示响应时发生错误: {str(e)}")
            self.send_button.setEnabled(True)
            self.send_button.setText("发送")

    def handle_edit_message(self, message):
        """处理编辑消息"""
        self.input_manager.setText(message)
        self.input_manager.setFocus()
        
    def closeEvent(self, event):
        """重写关闭事件"""
        # 保存当前的模型和对话状态
        current_model = self.model_selector.currentText()
        current_conversation_id = self.conversation_selector.currentData()
        
        # 将状态保存到assistant对象中
        self.assistant.last_model = current_model
        self.assistant.last_conversation_id = current_conversation_id
        
        event.ignore()
        self.assistant.show()
        self.hide()
        self.assistant.ai_main_window = None 

    def open_model_manager(self):
        """打开模型管理窗口"""
        model_manager = ModelManagerDialog(self.openai_chat, self)
        model_manager.model_updated.connect(self.refresh_models)
        model_manager.show()

    def refresh_models(self):
        """刷新模型列表并选中当前模型"""
        current_model = self.model_selector.currentText()
        self.load_models(select_model=current_model)

    def load_models(self, select_model=None):
        """加载模型列表"""
        models = self.openai_chat.get_all_models()
        self.model_selector.clear()
        for model_name in models:
            self.model_selector.addItem(model_name)
            
        if select_model and select_model in models:
            index = self.model_selector.findText(select_model)
            if index >= 0:
                self.model_selector.setCurrentIndex(index)

    def open_conversation_manager(self):
        """打开对话管理窗口"""
        current_model = self.model_selector.currentText()
        if not current_model:
            self.chat_display.display_error("请先选择一个模型！")
            return
            
        dialog = ConversationManagerDialog(self.conversation_manager, current_model, self)
        dialog.conversation_updated.connect(self.refresh_conversations)
        dialog.exec()
        
    def refresh_conversations(self):
        """刷新对话列表并选中当前对话"""
        current_model = self.model_selector.currentText()
        current_conversation_id = self.conversation_manager.current_conversation.get(current_model)
        self.update_conversation_selector(current_model, current_conversation_id)

    def on_model_changed(self, model_name):
        """处理模型变更"""
        self.load_models(select_model=model_name)

    def show_about(self):
        """显示关于信息"""
        QMessageBox.about(
            self,
            "关于",
            f"""<h3>AI助手 v{self.version}</h3>
            <p>一个本地的AI助手工具，支持多个AI模型的管理和对话。</p>
            <p><b>主要功能：</b></p>
            <ul>
                <li>支持多个AI模型的管理和切换</li>
                <li>对话历史的保存和管理</li>
                <li>实时响应和代码格式化</li>
                <li>便捷的悬浮窗操作</li>
            </ul>
            <p>作者：夜不能寐</p>
            <p>联系方式V+：X1012537840</p>
            <p>© 2024 保留所有权利</p>"""
        )